import argparse
import matplotlib.pyplot as plt
import mmcv
import os
from mmcv import Config
from pathlib import Path

from mmdet.core.utils import mask2ndarray
from mmdet.core.visualization import imshow_det_bboxes
from mmdet.datasets.builder import build_dataset


def parse_args():
    parser = argparse.ArgumentParser(description='Browse a dataset')
    parser.add_argument('--config', default='configs/_custom/faster_rcnn.py', help='train config file path')
    parser.add_argument(
        '--skip-type',
        type=str,
        nargs='+',
        default=['DefaultFormatBundle', 'Normalize', 'Collect'],
        help='skip some useless pipeline')
    parser.add_argument(
        '--output-dir',
        default=None,
        type=str,
        help='If there is no display interface, you can save it')
    parser.add_argument('--not-show', default=False, action='store_true')
    parser.add_argument(
        '--show-interval',
        type=float,
        default=0,
        help='the interval of show (s)')
    args = parser.parse_args()
    return args


def retrieve_data_cfg(config_path, skip_type):
    cfg = Config.fromfile(config_path)
    train_data_cfg = cfg.data.train
    train_data_cfg['pipeline'] = [
        x for x in train_data_cfg.pipeline if x['type'] not in skip_type
    ]

    return cfg


def main():
    args = parse_args()
    cfg = retrieve_data_cfg(args.config, args.skip_type)

    dataset = build_dataset(cfg.data.train)

    progress_bar = mmcv.ProgressBar(len(dataset))

    labels = ['liewen', 'bengque', 'duocengmian', 'yichang', 'wuzi', 'daojiaoyichang']
    xs = []
    ys = []
    for item in dataset:

        gt_bboxes = item.get('gt_bboxes', None)
        gt_labels = item.get('gt_labels', None)
        
        for idx, box in enumerate(gt_bboxes):
            label = labels[gt_labels[idx]]
            w = box[2] - box[0]
            h = box[3] - box[1]
            xs.append(w)
            ys.append(h)
        
        progress_bar.update()
    #
    plt.figure(figsize=(40, 40), dpi=80)
    plt.scatter(xs,ys)
    # plt.xticks(x)
    # plt.yticks(y)
    plt.xlabel('w')
    plt.ylabel('h')
    #plt.savefig("./label分部.png")
    plt.show()


if __name__ == '__main__':
    main()
